#!/usr/bin/env python
# -*- encoding: utf-8 -*-
import os
import time
import gradio as gr
from shuziren2d import txt2img, txt2video, audio2img

def video_identity(video):
    return video

## UI设计 ##
with gr.Blocks() as shuziren2D:
    # 顶部文字
    gr.Markdown("""
    # 2D数字人
    ### 开源项目地址：https://gitee.com/xingruispace/shuziren2d
    ### 需要优化声音等、输出视频质量等，可以在微信公众号星瑞元宇讨论
    """)

    demo_img = os.path.join(os.path.join(os.path.dirname(__file__),"file/demo/liqinzhao.jpg"))
    demo_txt = "常记溪亭日暮，沉醉不知归路。兴尽晚回舟，误入藕花深处。争渡，争渡，惊起一滩鸥鹭。"
    demo_video = os.path.join(os.path.join(os.path.dirname(__file__),"file/demo/liqinzhao_demo.mp4"))

    # # 顶部文字和选择，预留比如输入声音等
    # with gr.Row():
    #     with gr.Column():
    #         model_name = gr.Dropdown(utils.model_name_list, label="请选择一个基础模型", default="Baitian/momocha", multiselect=False)    # 同时作为训练lora的预训练模型
    # 多个tab
    with gr.Tabs():
        with gr.TabItem("文+图生成数字人"):
            # 一行 两列 左边一列是输入 右边一列是输出
            with gr.Row():
                with gr.Column():  # 左边一列是输入

                    txt2img_prompt = gr.Textbox(label="请输入需要播放的文字", lines=3, placeholder="请输入需要播放的文字", interactive=True, value=demo_txt)     

                    with gr.Row():
                        txt2img_img = gr.Image(label="请上传图片",value = demo_img)
   
                    # 生成、重置按钮（row：行）
                    with gr.Row():
                        txt2img_button = gr.Button("生成,预计需要3分钟")
                        # clear_button = gr.Button("重置")
                with gr.Column():  # 右边一列是输出
                    # 输出框
                    txt2img_output = gr.PlayableVideo(label='虚拟主播',value=demo_video )

        with gr.TabItem("文+视频生成数字人"):
            # 一行 两列 左边一列是输入 右边一列是输出
            with gr.Row():
                with gr.Column():  # 左边一列是输入

                    txt2video_prompt = gr.Textbox(label="请输入需要播放的文字", lines=3, placeholder="请输入需要播放的文字", interactive=True, value=demo_txt)     

                    with gr.Row():
                        txt2video_video = gr.Video(label="上传参考视频",value = demo_video)
   
                    # 生成、重置按钮（row：行）
                    with gr.Row():
                        txt2video_button = gr.Button("生成,预计需要3分钟")
                        # clear_button = gr.Button("重置")
                with gr.Column():  # 右边一列是输出
                    # 输出框
                    txt2video_output = gr.PlayableVideo(label='当作迁移到虚拟主播' )

        with gr.TabItem("数字人唱歌"):
            # 一行 两列 左边一列是输入 右边一列是输出
            with gr.Row():
                with gr.Column():  # 左边一列是输入

                    audio2img_Audio = gr.Audio(label="请输入需要唱的声音", interactive=True)
                    
                    with gr.Row():
                        audio2img_img = gr.Image(label="请上传图片",value = demo_img)
   
                    # 生成、重置按钮（row：行）
                    with gr.Row():
                        audio2img_button = gr.Button("生成,预计需要3分钟")
                        # clear_button = gr.Button("重置")
                with gr.Column():  # 右边一列是输出
                    # 输出框
                    audio2img_output = gr.PlayableVideo(label='虚拟主播',value=demo_video )
    
    # clear_button.click(cancels=txt2img_button.click)
    txt2img_button.click(fn=txt2img, inputs=[txt2img_prompt, txt2img_img],outputs=txt2img_output,api_name="txt2imgAPI",show_progress=True,queue=True)
    txt2video_button.click(fn=txt2video,inputs=[txt2video_prompt, txt2video_video],outputs=txt2video_output,show_progress=True)
    audio2img_button.click(fn=audio2img, inputs=[audio2img_Audio, audio2img_img],outputs=audio2img_output,show_progress=True,queue=True)
   


